package com.gitee.wsl.security.encrypt.sample

import com.gitee.wsl.ext.string.stripAllSpace
import com.gitee.wsl.security.encrypt.Encrypt
import com.gitee.wsl.security.encrypt.StringToStringEncrypt


/** @link https://wtool.com.cn/polybius.html */

class PolybiusSquareCipher(
   val table: String = TABLE_A_Z_WO_J,
   val encodeMap: String = DEFAULT_POLYBIUS_ENCODE_MAP,
   val replacePair: Pair<String, String> = "J" to "I"
) : StringToStringEncrypt {
    override fun decrypt(src: String): String {
        val properTable = table.stripAllSpace().uppercase()
        val map =
            properTable.associateBy {
                val i = properTable.indexOf(it)
                "${this.encodeMap[i / this.encodeMap.length]}${this.encodeMap[i % this.encodeMap.length]}"
            }
        val sb = StringBuilder()
        val tmp = StringBuilder()

        for (c in src) {
            when {
                c.isDigit() || c.isLetter() ->
                    if (tmp.length == 1) {
                        tmp.append(c)
                        sb.append(map[tmp.toString()] ?: tmp.toString())
                        tmp.clear()
                    } else {
                        tmp.append(c)
                    }
                else -> sb.append(c)
            }
        }
        return sb.toString()
    }

    override fun encrypt(src: String): String {
        val properTable = table.stripAllSpace().uppercase()
        val map =
            properTable.associateWith {
                val i = properTable.indexOf(it)
                "${encodeMap[i / encodeMap.length]}${encodeMap[i % encodeMap.length]}"
            }
        return src.uppercase()
            .replace(replacePair.first, replacePair.second)
            .map { map[it] ?: it }
            .joinToString("")
    }

    companion object{
        const val DEFAULT_POLYBIUS_ENCODE_MAP = "12345"
        val instance = PolybiusSquareCipher()
    }

}

val Encrypt.Companion.polybiusSquare get() = PolybiusSquareCipher.instance

fun Encrypt.Companion.polybiusSquare( table: String = TABLE_A_Z_WO_J,
                            encodeMap: String = PolybiusSquareCipher.DEFAULT_POLYBIUS_ENCODE_MAP,
                            replacePair: Pair<String, String> = "J" to "I"
) = PolybiusSquareCipher(table, encodeMap, replacePair)